package hbbft

import (
	"fmt"
	"math/big"

	"go.dedis.ch/kyber/v3"
	"go.dedis.ch/kyber/v3/group/edwards25519"
	"go.dedis.ch/kyber/v3/share"
	"go.dedis.ch/kyber/v3/util/random"
)

var suite = edwards25519.NewBlakeSHA256Ed25519()

// realCryptoProvider 基于kyber实现真实的TPKE和TBLS
type realCryptoProvider struct {
	// TPKE相关
	tpkeSuite   kyber.Group
	tpkePublic  *share.PubPoly  // 阈值加密公钥多项式
	tpkeShares  []*share.PriShare // 各节点的TPKE私钥份额
	tpkeN, tpkeT int

	// TBLS相关
	tblsSuite   kyber.Group
	tblsPublic  *share.PubPoly  // 阈值签名公钥多项式
	tblsShares  []*share.PriShare // 各节点的TBLS私钥份额
	tblsN, tblsT int
}

func (r *realCryptoProvider) TPKEPublic() interface{} {
	return r.tpkePublic
}

func (r *realCryptoProvider) TPKEPrivateFor(i int) interface{} {
	if i < 0 || i >= len(r.tpkeShares) {
		return nil
	}
	return r.tpkeShares[i]
}

func (r *realCryptoProvider) TBLSPublic() interface{} {
	return r.tblsPublic
}

func (r *realCryptoProvider) TBLSPrivateFor(i int) interface{} {
	if i < 0 || i >= len(r.tblsShares) {
		return nil
	}
	return r.tblsShares[i]
}

// realDKG 使用kyber的Pedersen DKG协议生成门限密钥
type realDKG struct{}

func (realDKG) Setup(N, T int) (CryptoProvider, error) {
	if T > N {
		return nil, fmt.Errorf("threshold T=%d cannot be greater than N=%d", T, N)
	}

	// 使用kyber的Pedersen DKG生成TPKE密钥
	tpkeShares, tpkePubPoly, err := generateThresholdKeys(suite, N, T)
	if err != nil {
		return nil, fmt.Errorf("TPKE DKG failed: %w", err)
	}

	// 使用kyber的Pedersen DKG生成TBLS密钥
	tblsShares, tblsPubPoly, err := generateThresholdKeys(suite, N, T)
	if err != nil {
		return nil, fmt.Errorf("TBLS DKG failed: %w", err)
	}

	return &realCryptoProvider{
		tpkeSuite:  suite,
		tpkePublic: tpkePubPoly,
		tpkeShares: tpkeShares,
		tpkeN:      N,
		tpkeT:      T,
		tblsSuite:  suite,
		tblsPublic: tblsPubPoly,
		tblsShares: tblsShares,
		tblsN:      N,
		tblsT:      T,
	}, nil
}

// generateThresholdKeys 使用简化版DKG生成阈值密钥
// 注意：这是简化的Trusted Dealer模式，真实分布式DKG需要多轮网络交互
func generateThresholdKeys(suite kyber.Group, N, T int) ([]*share.PriShare, *share.PubPoly, error) {
	// 创建秘密多项式：f(x) = secret + a1*x + ... + a_{T-1}*x^{T-1}
	secret := suite.Scalar().Pick(random.New())
	
	// 生成多项式系数（秘密值作为常数项）
	coefficients := make([]kyber.Scalar, T)
	coefficients[0] = secret
	for i := 1; i < T; i++ {
		coefficients[i] = suite.Scalar().Pick(random.New())
	}

	// 计算各节点的私钥份额 f(i+1)，i从0到N-1
	shares := make([]*share.PriShare, N)
	for i := 0; i < N; i++ {
		x := big.NewInt(int64(i + 1))
		shareValue := evaluatePolynomialKyber(suite, coefficients, x)
		shares[i] = &share.PriShare{
			I: i,
			V: shareValue,
		}
	}

	// 计算公钥多项式：G^f(x) 的系数
	publicCoeffs := make([]kyber.Point, T)
	for i := 0; i < T; i++ {
		publicCoeffs[i] = suite.Point().Mul(coefficients[i], nil)
	}
	
	pubPoly := share.NewPubPoly(suite, suite.Point().Base(), publicCoeffs)

	return shares, pubPoly, nil
}

// evaluatePolynomialKyber 在点x处求值多项式 f(x) = sum(coeffs[i] * x^i)
func evaluatePolynomialKyber(suite kyber.Group, coefficients []kyber.Scalar, x *big.Int) kyber.Scalar {
	result := suite.Scalar().Zero()
	xx := big.NewInt(1)
	
	for _, coeff := range coefficients {
		xScalar := suite.Scalar().SetInt64(xx.Int64())
		term := suite.Scalar().Mul(coeff, xScalar)
		result = suite.Scalar().Add(result, term)
		xx.Mul(xx, x)
	}
	
	return result
}
